/*
 * LOH.java
 *
 * Copyright (C) 2002-2006 Alexei Drummond and Andrew Rambaut
 *
 * This file is part of BEAST.
 * See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership and licensing.
 *
 * BEAST is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 *  BEAST is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with BEAST; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301  USA
 */

package dr.evomodel.substmodel;

import dr.inference.model.Model;
import dr.inference.model.Parameter;
//import dr.inference.model.Statistic;
import dr.xml.*;
import dr.evolution.datatype.*;

import java.util.logging.Logger;

/**
 * LOH model 
 *
 * @version $Id: LOH.java $
 *
 * @author Rumen Kostadinov
 */
public class LOH extends AbstractSubstitutionModel
{
    public static final String LOH_MODEL = "lohModel";
    public static final String ALPHA = "alpha";
    public static final String BETA = "beta";
    public static final String FREQUENCIES = "frequencies";
    
    private Parameter alphaParameter = null;
    private Parameter betaParameter = null;

    
    private boolean updateIntermediates = true;

    /** Used for precalculations */
    //private double 				beta, A_R, A_Y;
    //private double 				tab1A, tab2A, tab3A;
    //private double 				tab1C, tab2C, tab3C;
    //private double 				tab1G, tab2G, tab3G;
    //private double 				tab1T, tab2T, tab3T;

    /**
     * Constructor
     */
    public LOH(Parameter alphaParameter, Parameter betaParameter, DataType dataType, FrequencyModel freqModel) {
    	super(LOH_MODEL, dataType, freqModel);
        this.alphaParameter = alphaParameter;
        this.betaParameter = betaParameter;
        
        addParameter(alphaParameter);
        addParameter(betaParameter);
        
        alphaParameter.addBounds(new Parameter.DefaultBounds(Double.POSITIVE_INFINITY, 0.0, 1));
        betaParameter.addBounds(new Parameter.DefaultBounds(Double.POSITIVE_INFINITY, 0.0, 1));
        
        updateIntermediates = true;

        //addStatistic(tsTvStatistic);
    }

    /**
     * set alpha
     */
    public void setAlpha(double alpha)
    {
        alphaParameter.setParameterValue(0, alpha);
        updateMatrix = true;
    }
    
    /**
     * set beta
     */
    public void setBeta(double beta)
    {
        betaParameter.setParameterValue(0, beta);
        updateMatrix = true;
    }
    
    /**
     * @return alpha
     */
    public final double getAlpha() { return alphaParameter.getParameterValue(0); }

    /**
     * @return beta
     */
    public final double getBeta() { return betaParameter.getParameterValue(0); }

    /**
     * set ts/tv
     */
    //public void setTsTv(double tsTv) {
    //    this.tsTv = tsTv;
    //      calculateFreqRY();
    //    setKappa((tsTv*freqR*freqY)/(freqA*freqG+freqC*freqT));
    //}

    /**
     * @return tsTv
     */
    //public double getTsTv()
    //{
    //      calculateFreqRY();
     //     tsTv = (getKappa() * (freqA*freqG + freqC*freqT))/(freqR*freqY);

     //   return tsTv;
    //}

    protected void frequenciesChanged() {
        // frequencyModel changed
        updateIntermediates = true;
    }

    protected void ratesChanged() {
		// Nothing to precalculate
	}
    
    private void calculateIntermediates() {

        //calculateFreqRY();
/*
        tab1A = freqA*((1/freqR)-1);
        tab2A = (freqR-freqA)/freqR;
        tab3A = freqA/freqR;
        tab1C = freqC*((1/freqY)-1);
        tab2C = (freqY-freqC)/freqY;
        tab3C = freqC/freqY;
        tab1G = freqG*((1/freqR)-1);
        tab2G = (freqR-freqG)/freqR;
        tab3G = freqG/freqR;
        tab1T = freqT*((1/freqY)-1);
        tab2T = (freqY-freqT)/freqY;
        tab3T = freqT/freqY;
*/
        updateMatrix = true;
        updateIntermediates = false;
 
    }

    /**
     * get the complete transition probability matrix for the given distance
     *
     * @param distance the expected number of substitutions
     * @param matrix an array to store the matrix
     */
    public void getTransitionProbabilities(double distance, double[] matrix)
    {
        if (updateIntermediates) {
            calculateIntermediates();
        }

        if (updateMatrix) {
            setupMatrix();
        }
        
        // LOH Model matrix
        // from/to       0     1
        //         0    1-a    a    
        //         1     b    1-b
        
        double alpha = getAlpha();
        double beta = getBeta();
        
        matrix[1] = alpha * (1-Math.exp(-(alpha+beta)*distance)) / (alpha+beta);
        matrix[2] = beta  * (1-Math.exp(-(alpha+beta)*distance)) / (alpha+beta);
        matrix[0] = 1 - matrix[1];
        matrix[3] = 1 - matrix[2];
        
/*
        double xx, aa, bbR, bbY;

        xx = beta * distance;
        bbR = Math.exp(xx*A_R);
        bbY = Math.exp(xx*A_Y);

        aa = Math.exp(xx);

        matrix[0] =	freqA+(tab1A*aa)+(tab2A*bbR);
        matrix[1] =	freqC*(1-aa);
        matrix[2] =	freqG+(tab1G*aa)-(tab3G*bbR);
        matrix[3] =	freqT*(1-aa);

        matrix[4] =	freqA*(1-aa);
        matrix[5] =	freqC+(tab1C*aa)+(tab2C*bbY);
        matrix[6] =	freqG*(1-aa);
        matrix[7] =	freqT+(tab1T*aa)-(tab3T*bbY);

        matrix[8] =	freqA+(tab1A*aa)-(tab3A*bbR);
        matrix[9] =	matrix[1];
        matrix[10] =freqG+(tab1G*aa)+(tab2G*bbR);
        matrix[11] =matrix[3];

        matrix[12] =matrix[4];
        matrix[13] =freqC+(tab1C*aa)-(tab3C*bbY);
        matrix[14] =matrix[6];
        matrix[15] =freqT+(tab1T*aa)+(tab2T*bbY);

        C = new double[N][N];
        for (int i = 0; i < N; i++)
            for (int j = 0; j < N; j++)
                for (int k = 0; k < N; k++)
                    C[i][k] += A[i][j] * B[j][k];
*/
    }

    /**
     * setup substitution matrix
     */
    protected void setupMatrix()
    {
    	/*
        double kappa = getKappa();
        beta = -1.0 / (2.0*(freqR*freqY + kappa*(freqA*freqG +
                                                    freqC*freqT)));

        A_R = 1.0 +freqR*(kappa-1);
        A_Y = 1.0 +freqY*(kappa-1);
        */
    	
        updateMatrix = false;
    }

    protected void setupRelativeRates() { }

    // *****************************************************************
    // Interface Model
    // *****************************************************************

    protected void storeState() { } // nothing to do
    
    /**
     * Restore the stored state
     */
    public void restoreState() {
        //super.restoreState();
        updateIntermediates = true;
        updateMatrix = true;
    }

    /**
     * Restore the stored state
     */
    public void adoptState(Model source) {
        //super.adoptState(source);
        updateIntermediates = true;
        updateMatrix = true;
    }

    // **************************************************************
    // XHTMLable IMPLEMENTATION
    // **************************************************************

    public String toXHTML() {
        StringBuffer buffer = new StringBuffer();

       /* buffer.append("<em>HKY Model</em> Ts/Tv = ");
        buffer.append(getTsTv());
        buffer.append(" (kappa = ");
        buffer.append(getKappa());
        buffer.append(")");
        */
        return buffer.toString();
    }

    /**
     * Parses an element from an DOM document into a DemographicModel. Recognises
     * ConstantPopulation and ExponentialGrowth.
     */
    public static XMLObjectParser PARSER = new AbstractXMLObjectParser() {

        public String getParserName() { return LOH_MODEL; }

        public Object parseXMLObject(XMLObject xo) throws XMLParseException {

            //Parameter kappaParam = (Parameter)xo.getSocketChild(KAPPA);
            
            Parameter alphaParam = (Parameter)xo.getSocketChild(ALPHA);
            Parameter betaParam = (Parameter)xo.getSocketChild(BETA);
            
            //FrequencyModel freqModel = (FrequencyModel)xo.getSocketChild(FREQUENCIES);
            XMLObject cxo = (XMLObject)xo.getChild(FREQUENCIES);
			FrequencyModel freqModel = (FrequencyModel)cxo.getChild(FrequencyModel.class);
			DataType dataType = null;
			
			if (xo.hasAttribute(DataType.DATA_TYPE)) {
				String dataTypeStr = xo.getStringAttribute(DataType.DATA_TYPE);
				if (dataTypeStr.equals(Nucleotides.DESCRIPTION)) {
					dataType = Nucleotides.INSTANCE;
				} else if (dataTypeStr.equals(AminoAcids.DESCRIPTION)) {
					dataType = AminoAcids.INSTANCE;
				} else if (dataTypeStr.equals(Codons.DESCRIPTION)) {
					dataType = Codons.UNIVERSAL;
				} else if (dataTypeStr.equals(TwoStates.DESCRIPTION)) {
					dataType = TwoStates.INSTANCE;
				}
			}
							
			if (dataType == null) dataType = (DataType)xo.getChild(DataType.class);
			
            Logger.getLogger("dr.evomodel").info("Creating LOH substitution model. Initial alpha = " 
            		+ alphaParam.getParameterValue(0) + " beta = " + betaParam.getParameterValue(0));

            return new LOH(alphaParam, betaParam, dataType, freqModel);
        }

        //************************************************************************
        // AbstractXMLObjectParser implementation
        //************************************************************************

        public String getParserDescription() {
            return "This element represents an instance of the LOH (Felsenstein, Kuhner, Maley, Kostadinov 2007) model of LOH evolution.";
        }

        public Class getReturnType() { return LOH.class; }

        public XMLSyntaxRule[] getSyntaxRules() { return rules; }

        private XMLSyntaxRule[] rules = new XMLSyntaxRule[] {
            //new ElementRule(FREQUENCIES,
            //    new XMLSyntaxRule[] { new ElementRule(FrequencyModel.class) }),
        	new XORRule(
        				new StringAttributeRule(DataType.DATA_TYPE, "The type of sequence data", new String[] { Nucleotides.DESCRIPTION, AminoAcids.DESCRIPTION, Codons.DESCRIPTION, TwoStates.DESCRIPTION}, false),
        				new ElementRule(DataType.class)
        				),
        	new ElementRule(FREQUENCIES, FrequencyModel.class),
            new ElementRule(ALPHA,
                new XMLSyntaxRule[] { new ElementRule(Parameter.class) }),
            new ElementRule(BETA,
                new XMLSyntaxRule[] { new ElementRule(Parameter.class) })
        };

    };

    //
    // Private stuff
    //
/*
    private Statistic tsTvStatistic = new Statistic.Abstract() {

        public String getStatisticName() {
            return "tsTv";
        }

        public int getDimension() { return 1; }

        public double getStatisticValue(int dim) {
            return getTsTv();
        }

    };
*/
    
}
